import concurrent
import os
import threading
from concurrent.futures import ThreadPoolExecutor

import snowflake.connector
from PySide6.QtWidgets import QDialog, QMessageBox

from utils.signal_emitter import SignalEmitter
from windows.load_to_db_dialog_ui import Ui_LoadToRemoteDBDialog


class LoadToRemoteDBDialog(QDialog):
    def __init__(self, parent=None):
        super(LoadToRemoteDBDialog, self).__init__(parent)

        self.table_name = parent.table_name

        self.all_processed_count = 0
        self.parent = parent
        self.db_dialog_ui = Ui_LoadToRemoteDBDialog()
        # 设置dialog的高
        self.db_dialog_ui.setupUi(self)
        QDialog.resize(self, 550, 420)
        QDialog.setWindowTitle(self, f"Upload {self.table_name} To DB")

        db_source = self.db_dialog_ui.dbSourcecomboBox.currentText()
        self.db_dialog_ui.sourceFileLabel.hide()
        self.db_dialog_ui.sourceFileLineEdit.hide()
        self.db_dialog_ui.sourceQueryLabel.hide()
        self.db_dialog_ui.sourceQueryLineEdit.hide()
        self.db_dialog_ui.targetTableLabel.hide()
        self.db_dialog_ui.targetTableLineEdit.hide()
        self.db_dialog_ui.submitButton.hide()
        self.db_dialog_ui.bacckButton.hide()
        self.db_dialog_ui.closeButton.hide()
        self.db_dialog_ui.loadProgressBar.hide()
        self.db_dialog_ui.loadedCuntLabel.hide()
        self.db_dialog_ui.threadCountLabel.hide()
        self.db_dialog_ui.threadCountSpinBox.hide()

        self.db_dialog_ui.sourceQueryLineEdit.setEnabled(False)
        self.db_dialog_ui.sourceFileLineEdit.setEnabled(False)

        self.db_dialog_ui.nextButton.clicked.connect(lambda: self.next_action())
        self.db_dialog_ui.bacckButton.clicked.connect(lambda: self.back_action())
        self.db_dialog_ui.submitButton.clicked.connect(lambda: self.submit_action())
        self.db_dialog_ui.closeButton.clicked.connect(lambda: self.close())

        self.db_conn = None

        if db_source == 'Snowflake':
            self.db_dialog_ui.passwordLabel.hide()
            self.db_dialog_ui.passwordLineEdit.hide()
            self.db_dialog_ui.portLabel.hide()
            self.db_dialog_ui.portLineEdit.hide()

            self.db_dialog_ui.hostLineEdit.setText('lz64309.east-us-2.azure.snowflakecomputing.com')
            username = os.getlogin()
            self.db_dialog_ui.userLineEdit.setText(username + '@jci.com')
            self.db_dialog_ui.roleLineEdit.setText('SF_BUILD_RO_USABLE_TRANS')
            self.db_dialog_ui.warehouseLineEdit.setText('EDP_BUILD_RO_WH')

    def start_load_to_db_thread(self, signal_emitter: SignalEmitter, source_query: str, target_table: str,
                                thread_count: int):
        self.all_processed_count = 0
        try:
            self.snowflake_connection()
        except Exception as e:
            signal_emitter.result_signal.emit(False, 'Error: ' + str(e), None)

        signal_emitter.result_signal.emit(True, "0", f"Start loading data to Snowflake...")
        if self.parent.file_type.lower() != 'csv':
            source_query = source_query.replace(self.table_name, self.parent.file_path)
        source_cursor = self.parent.conn.cursor()
        target_cursor = self.db_conn.cursor()

        try:
            source_duckdb = source_cursor.sql(source_query)
            source_count = source_duckdb.shape[0]
            source_columns = source_duckdb.columns
            # 每次查询的数据量
            batch_size = 10000

            # target_table_info = target_cursor.execute("DESC TABLE " + target_table)
            # target_table_infor_data = target_table_info.fetchall()
            placeholders = ", ".join(["%s" for _ in source_columns])
            source_columns = ['"' + col + '"' for col in source_columns]
            target_insert_query = f"INSERT INTO {target_table} ({', '.join(source_columns).upper()}) VALUES ({placeholders})"
            target_cursor.execute("BEGIN;")
            # 创建锁对象用于保证线程安全
            lock = threading.Lock()

            def load_to_db(default_offset=0, default_count=0):
                sub_source_cursor = self.parent.conn.cursor()
                current_thread_id = threading.current_thread().ident
                try:
                    for offset in range(default_offset, default_count, batch_size):
                        this_batch_size = batch_size
                        if default_count < offset + batch_size:
                            this_batch_size = default_count - offset
                        sql_query = f"{source_query} LIMIT {this_batch_size} OFFSET {offset};"
                        data = sub_source_cursor.sql(sql_query).fetchall()
                        print(
                            f"当前线程ID: {current_thread_id}, 当前Offset: {offset}, "
                            f"当前Count: {default_count}, 当前Len: {len(data)}")
                        target_cursor.executemany(target_insert_query, data)
                        # 计算已处理数据量的百分比
                        process_count = this_batch_size if offset + this_batch_size < default_count else default_count - offset
                        lock.acquire()  # 获取锁，保证同一时刻只有一个线程能操作共享变量
                        self.all_processed_count += process_count
                        percentage = int(self.all_processed_count / source_count * 100)
                        signal_emitter.result_signal.emit(True, str(percentage),
                                                          f"{self.all_processed_count:,}/{source_count:,}")
                        lock.release()
                    return True, ""
                except Exception as e:
                    print(f"Error occurred in thread {current_thread_id}: {str(e)}")
                    return False, f"Error occurred in thread {current_thread_id}: {str(e)}"
                finally:
                    sub_source_cursor.close()

            with concurrent.futures.ThreadPoolExecutor(max_workers=thread_count) as t:  # 创建一个最大容纳数量为10的线程池
                total_num = source_count
                base_size = total_num // thread_count
                remainder = total_num % thread_count
                start_value = 0
                futures = []
                total_count = 0
                for i in range(thread_count):
                    current_size = base_size
                    if i + 1 == thread_count:
                        current_size += remainder
                    total_count += current_size
                    print(f"第{i + 1}份，开始值为{start_value}，大小为{total_count}")
                    future = t.submit(load_to_db, start_value, total_count)
                    futures.append(future)
                    start_value += current_size
                # 等待所有任务完成
                done, no_done = concurrent.futures.wait(futures)

                flag = True
                error_msg = ''
                for future in done:
                    status, e = future.result()
                    if not status:
                        error_msg += f"{future.result()}\n"
                        flag = False

                if flag:
                    self.db_conn.commit()
                    signal_emitter.result_signal.emit(True, 'success', None)
                else:
                    signal_emitter.result_signal.emit(False, error_msg, None)

        except Exception as e:
            print(e)
            self.db_conn.rollback()
            signal_emitter.result_signal.emit(False, 'Error: ' + str(e), None)
        finally:
            target_cursor.close()
            source_cursor.close()

    def load_progress_signal(self, value, params1, params2):
        if value:
            if params1 == 'success':
                QMessageBox.information(self, "Success", f"{self.table_name} upload successful!")
                self.db_dialog_ui.submitButton.setEnabled(True)
                self.db_dialog_ui.bacckButton.setEnabled(True)
                self.db_dialog_ui.targetTableLineEdit.setEnabled(True)
                self.db_dialog_ui.threadCountSpinBox.setEnabled(True)
                self.db_dialog_ui.submitButton.hide()
                self.db_dialog_ui.bacckButton.hide()
                self.db_dialog_ui.closeButton.show()
                return
            self.db_dialog_ui.loadProgressBar.setValue(int(params1))
            self.db_dialog_ui.loadedCuntLabel.setText(str(params2))
        else:
            QMessageBox.critical(self, "Error", f"{self.table_name} upload failed!\n{params1}")
            self.db_dialog_ui.submitButton.setEnabled(True)
            self.db_dialog_ui.bacckButton.setEnabled(True)
            self.db_dialog_ui.loadProgressBar.hide()
            self.db_dialog_ui.loadedCuntLabel.hide()
            self.db_dialog_ui.targetTableLineEdit.setEnabled(True)
            self.db_dialog_ui.threadCountSpinBox.setEnabled(True)

    def submit_action(self):

        target_table_text = self.db_dialog_ui.targetTableLineEdit.text()

        if target_table_text == '':
            QMessageBox.warning(self, "Warning", "Target table name cannot be empty!")
            return

        self.db_dialog_ui.loadedCuntLabel.setText('Logging in to database...')

        self.db_dialog_ui.loadProgressBar.show()
        self.db_dialog_ui.loadedCuntLabel.show()
        self.db_dialog_ui.targetTableLineEdit.setEnabled(False)
        self.db_dialog_ui.threadCountSpinBox.setEnabled(False)

        self.db_dialog_ui.submitButton.setEnabled(False)
        self.db_dialog_ui.bacckButton.setEnabled(False)

        source_query = self.db_dialog_ui.sourceQueryLineEdit.text()
        target_table = self.db_dialog_ui.targetTableLineEdit.text()
        thread_count = self.db_dialog_ui.threadCountSpinBox.value()
        signal_emitter = SignalEmitter()
        signal_emitter.result_signal.connect(self.load_progress_signal)
        submit_thread = threading.Thread(target=self.start_load_to_db_thread,
                                         args=(signal_emitter, source_query, target_table, thread_count))
        submit_thread.start()

    def next_action(self):

        self.db_dialog_ui.dbSourcecomboBox.hide()
        self.db_dialog_ui.dBSourceLabel.hide()
        self.db_dialog_ui.hostLabel.hide()
        self.db_dialog_ui.hostLineEdit.hide()
        self.db_dialog_ui.portLabel.hide()
        self.db_dialog_ui.portLineEdit.hide()
        self.db_dialog_ui.userLabel.hide()
        self.db_dialog_ui.userLineEdit.hide()
        self.db_dialog_ui.snowflakeRoleLabel.hide()
        self.db_dialog_ui.roleLineEdit.hide()
        self.db_dialog_ui.snowflakeWarehouseLabel.hide()
        self.db_dialog_ui.warehouseLineEdit.hide()
        self.db_dialog_ui.nextButton.hide()

        self.db_dialog_ui.sourceFileLabel.show()
        self.db_dialog_ui.sourceFileLineEdit.show()
        self.db_dialog_ui.sourceQueryLabel.show()
        self.db_dialog_ui.sourceQueryLineEdit.show()
        self.db_dialog_ui.targetTableLabel.show()
        self.db_dialog_ui.targetTableLineEdit.show()
        self.db_dialog_ui.submitButton.show()
        self.db_dialog_ui.bacckButton.show()
        self.db_dialog_ui.threadCountLabel.show()
        self.db_dialog_ui.threadCountSpinBox.show()

        file = self.parent.csv_reader_ui.filePathLineEdit.text()
        query = self.parent.csv_reader_ui.sQLLineEdit.text()

        query = query.split('LIMIT')[0].strip()
        self.db_dialog_ui.sourceFileLineEdit.setText(file)
        self.db_dialog_ui.sourceQueryLineEdit.setText(query)

    def closeEvent(self, event):
        # 在这里添加关闭时需要执行的操作
        print("窗口即将关闭")
        if self.db_conn is not None:
            self.db_conn.close()
        # 可以选择接受或忽略关闭事件
        event.accept()  # 接受关闭事件
        # event.ignore()  # 忽略关闭事件

    def back_action(self):
        self.db_dialog_ui.dbSourcecomboBox.show()
        self.db_dialog_ui.dBSourceLabel.show()
        self.db_dialog_ui.hostLabel.show()
        self.db_dialog_ui.hostLineEdit.show()
        self.db_dialog_ui.userLabel.show()
        self.db_dialog_ui.userLineEdit.show()
        self.db_dialog_ui.snowflakeRoleLabel.show()
        self.db_dialog_ui.roleLineEdit.show()
        self.db_dialog_ui.snowflakeWarehouseLabel.show()
        self.db_dialog_ui.warehouseLineEdit.show()
        self.db_dialog_ui.nextButton.show()

        self.db_dialog_ui.bacckButton.hide()
        self.db_dialog_ui.sourceFileLabel.hide()
        self.db_dialog_ui.sourceFileLineEdit.hide()
        self.db_dialog_ui.sourceQueryLabel.hide()
        self.db_dialog_ui.sourceQueryLineEdit.hide()
        self.db_dialog_ui.targetTableLabel.hide()
        self.db_dialog_ui.targetTableLineEdit.hide()
        self.db_dialog_ui.submitButton.hide()
        self.db_dialog_ui.threadCountLabel.hide()
        self.db_dialog_ui.threadCountSpinBox.hide()

        self.db_dialog_ui.loadProgressBar.hide()
        self.db_dialog_ui.loadedCuntLabel.hide()

    def snowflake_connection(self):
        host = self.db_dialog_ui.hostLineEdit.text()
        user = self.db_dialog_ui.userLineEdit.text()
        role = self.db_dialog_ui.roleLineEdit.text()
        warehouse = self.db_dialog_ui.warehouseLineEdit.text()
        account = host.split('.')[0]

        params = {
            'Authentication': 'authenticator',
            'authenticator': 'externalbrowser',
            'host': host,
            'role': role,
            'warehouse': warehouse,
            'user': user,
            'account': account
        }

        self.db_conn = snowflake.connector.connect(**params)